import numpy as np
import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation


class Animation:
    def __init__(self):
        self.N = None  # 轨迹的数量
        self.trajectory = []  # 存储轨迹的列表

    def load_trajectory(self, trajectory):
        """
        加载单条轨迹数据。
        轨迹是一个列表，列表中的每个元素为一个 np.array([x, y, v, h])。
        """
        self.trajectory.append(trajectory)

    def calculate_max_x_y(self):
        """
        计算所有轨迹的最大值和最小值，用于设置图表的边界。
        返回 x_min, x_max, y_min, y_max
        """
        x_min, x_max = float('inf'), float('-inf')
        y_min, y_max = float('inf'), float('-inf')

        # 遍历所有轨迹
        for traj in self.trajectory:
            x_data = [t[0] for t in traj]
            y_data = [t[1] for t in traj]

            # 更新 x 和 y 的最小最大值
            x_min = min(x_min, min(x_data))
            x_max = max(x_max, max(x_data))
            y_min = min(y_min, min(y_data))
            y_max = max(y_max, max(y_data))

        return x_min, x_max, y_min, y_max

    def draw(self):
        """
        绘制动画，每条轨迹将随着帧的变化逐步显示。
        """
        fig, ax = plt.subplots()

        # 计算边界
        x_min, x_max, y_min, y_max = self.calculate_max_x_y()
        ax.set_xlim(x_min - 10, x_max + 10)  # 稍微留一点边距
        ax.set_ylim(y_min - 10, y_max + 10)

        # 轨迹的数量
        lines = []
        colors = plt.cm.jet(np.linspace(0, 1, len(self.trajectory)))  # 使用 jet colormap 为每条轨迹分配颜色

        for i, traj in enumerate(self.trajectory):
            line, = ax.plot([], [], 'o-', markersize=8, color=colors[i])  # 每条轨迹的线，设置不同颜色
            lines.append(line)

        # 初始化函数
        def init():
            for line in lines:
                line.set_data([], [])  # 初始化为空
            return lines

        # 更新函数：每次调用时更新轨迹的点
        def update(frame):
            for idx, traj in enumerate(self.trajectory):
                # 获取当前轨迹的 x, y 数据，只显示当前点
                x_data = [traj[frame][0]]
                y_data = [traj[frame][1]]
                lines[idx].set_data(x_data, y_data)  # 仅更新当前点

            return lines

        # 创建动画
        ani = FuncAnimation(fig, update, frames=max(len(traj) for traj in self.trajectory),
                            init_func=init, blit=True, interval=1000)
        ani.save("test.gif")
        # 显示图表


# 示例使用
if __name__ == "__main__":
    # 创建动画对象
    anim = Animation()

    # 加载轨迹数据，每条轨迹是一个包含 np.array([x, y, v, h]) 的列表
    traj1 = [np.array([0, 0, 10, 0]), np.array([1, 1, 12, 1]), np.array([2, 4, 15, 2])]
    traj2 = [np.array([0, 0, 10, 0]), np.array([1, 2, 13, 1]), np.array([2, 5, 14, 2])]
    anim.load_trajectory(traj1)
    anim.load_trajectory(traj2)

    # 绘制动画
    anim.draw()